{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Unet_Implementation_PyTorch.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Qiuyan918/Unet_Implementation_PyTorch/blob/master/Unet_Implementation_PyTorch.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SJJcl5Day-aS",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import cv2\n",
        "from tqdm import tqdm\n",
        "from copy import deepcopy\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "import torchvision\n",
        "from torch.utils.data import Dataset\n",
        "from torch.utils.data import DataLoader\n",
        "from torch.autograd import Variable\n",
        "from torch.utils.data.sampler import RandomSampler"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E3BehnPezvFY",
        "colab_type": "text"
      },
      "source": [
        "# Configs"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YPEAHcCDzzQ6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "n_fold = 5\n",
        "pad_left = 27\n",
        "pad_right = 27\n",
        "fine_size = 202\n",
        "batch_size = 18\n",
        "epoch = 300 \n",
        "snapshot = 6 \n",
        "max_lr = 0.012 \n",
        "min_lr = 0.001 \n",
        "momentum = 0.9 \n",
        "weight_decay = 1e-4 \n",
        "n_fold = 5\n",
        "device = torch.device('cuda')\n",
        "save_weight = 'weights/'\n",
        "if not os.path.isdir(save_weight):\n",
        "  os.mkdir(save_weight)\n",
        "weight_name = 'model_' + str(fine_size+pad_left+pad_right) + '_res18' \n",
        "\n",
        "train_image_dir = 'tgs-salt-identification-challenge/train/images'\n",
        "train_mask_dir = 'tgs-salt-identification-challenge/train/masks'\n",
        "test_image_dir = 'tgs-salt-identification-challenge/test/images'"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KY7UTuce1d0T",
        "colab_type": "text"
      },
      "source": [
        "# Split"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v90Ul61N0hKF",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "depths = pd.read_csv('tgs-salt-identification-challenge/depths.csv')\n",
        "depths.sort_values('z', inplace=True)\n",
        "depths.drop('z', axis=1, inplace=True)\n",
        "depths['fold'] = (list(range(0,5)) * depths.shape[0])[:depths.shape[0]]\n",
        "\n",
        "train_df = pd.read_csv('tgs-salt-identification-challenge/train.csv')\n",
        "train_df = train_df.merge(depths)\n",
        "dist = []\n",
        "for id in train_df.id.values:\n",
        "  img = cv2.imread(f'tgs-salt-identification-challenge/train/images/{id}.png', cv2.IMREAD_GRAYSCALE)\n",
        "  dist.append(np.unique(img).shape[0])\n",
        "train_df['unique_pixels'] = dist"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UI0u3k-l564p",
        "colab_type": "text"
      },
      "source": [
        "# Dataset"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9MDKVsgDXZ5g",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def trainImageFetch(images_id):\n",
        "  image_train = np.zeros((images_id.shape[0], 101, 101), dtype=np.float32)\n",
        "  mask_train = np.zeros((images_id.shape[0], 101, 101), dtype=np.float32)\n",
        "\n",
        "  for idx, image_id in tqdm(enumerate(images_id), total=images_id.shape[0]):\n",
        "    image_path = os.path.join(train_image_dir, image_id+'.png')\n",
        "    mask_path = os.path.join(train_mask_dir, image_id+'.png')\n",
        "\n",
        "    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE).astype(np.float32) / 255\n",
        "    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE).astype(np.float32) / 255\n",
        "\n",
        "    image_train[idx] = image\n",
        "    mask_train[idx] = mask\n",
        "  \n",
        "  return image_train, mask_train\n",
        "\n",
        "def testImageFetch(test_id):\n",
        "  image_test = np.zeros((len(test_id), 101, 101), dtype=np.float32)\n",
        "\n",
        "  for idx, image_id in tqdm(enumerate(test_id), total=len(test_id)):\n",
        "    image_path = os.path.join(test_image_dir, image_id+'.png')\n",
        "    image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE).astype(np.float32) / 255\n",
        "    image_test[idx] = image\n",
        "\n",
        "  return image_test\n",
        "\n",
        "def do_resize2(image, mask, H, W):\n",
        "  image = cv2.resize(image, dsize=(W,H))\n",
        "  mask = cv2.resize(mask, dsize=(W,H))\n",
        "  return image, mask\n",
        "\n",
        "def do_center_pad(image, pad_left, pad_right):\n",
        "  return np.pad(image, (pad_left, pad_right), 'edge')\n",
        "\n",
        "def do_center_pad2(image, mask, pad_left, pad_right):\n",
        "  image = do_center_pad(image, pad_left, pad_right)\n",
        "  mask = do_center_pad(mask, pad_left, pad_right)\n",
        "  return image, mask\n",
        "\n",
        "class SaltDataset(Dataset):\n",
        "  def __init__(self, image_list, mode, mask_list=None, fine_size=202, pad_left=0, pad_right=0):\n",
        "    self.imagelist = image_list\n",
        "    self.mode = mode\n",
        "    self.masklist = mask_list\n",
        "    self.fine_size = fine_size\n",
        "    self.pad_left = pad_left\n",
        "    self.pad_right = pad_right\n",
        "\n",
        "  def __len__(self):\n",
        "    return len(self.imagelist)\n",
        "\n",
        "  def __getitem__(self, idx):\n",
        "    image = deepcopy(self.imagelist[idx])\n",
        "\n",
        "    if self.mode == 'train':\n",
        "      mask = deepcopy(self.masklist[idx])\n",
        "      label = np.where(mask.sum() == 0, 1.0, 0.0).astype(np.float32)\n",
        "\n",
        "      if self.fine_size != image.shape[0]:\n",
        "        image, mask = do_resize2(image, mask, self.fine_size, self.fine_size)\n",
        "\n",
        "      if self.pad_left != 0:\n",
        "        image, mask = do_center_pad2(image, mask, self.pad_left, self.pad_right)\n",
        "\n",
        "      image = image.reshape(1, image.shape[0], image.shape[1])\n",
        "      mask = mask.reshape(1, mask.shape[0], mask.shape[1])    \n",
        "\n",
        "      return image, mask, label\n",
        "\n",
        "    elif self.mode == 'val':\n",
        "      mask = deepcopy(self.masklist[idx])\n",
        "\n",
        "      if self.fine_size != image.shape[0]:\n",
        "        image, mask = do_resize2(image, mask, self.fine_size, self.fine_size)\n",
        "\n",
        "      if self.pad_left != 0:\n",
        "        image = do_center_pad(image, self.pad_left, self.pad_right)\n",
        "\n",
        "      image = image.reshape(1, image.shape[0], image.shape[1])\n",
        "      mask = mask.reshape(1, mask.shape[0], mask.shape[1])  \n",
        "\n",
        "      return image, mask\n",
        "\n",
        "    elif self.mode == 'test':\n",
        "      if self.fine_size != image.shape[0]:\n",
        "        image = cv2.resize(image, dsize=(self.fine_size, self.fine_size))\n",
        "\n",
        "      if self.pad_left != 0:\n",
        "        image = do_center_pad(image, self.pad_left, self.pad_right)\n",
        "\n",
        "      image = image.reshape(1, image.shape[0], image.shape[1])\n",
        "\n",
        "      return image     "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t5gnjD4X69f6",
        "colab_type": "text"
      },
      "source": [
        "# Model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hZz39glnTum1",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Decoder(nn.Module):\n",
        "  def __init__(self, in_channels, middle_channels, out_channels):\n",
        "    super(Decoder, self).__init__()\n",
        "    self.up = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)\n",
        "    self.conv_relu = nn.Sequential(\n",
        "        nn.Conv2d(middle_channels, out_channels, kernel_size=3, padding=1),\n",
        "        nn.ReLU(inplace=True)\n",
        "        )\n",
        "  def forward(self, x1, x2):\n",
        "    x1 = self.up(x1)\n",
        "    x1 = torch.cat((x1, x2), dim=1)\n",
        "    x1 = self.conv_relu(x1)\n",
        "    return x1\n",
        "\n",
        "class UNet(nn.Module):\n",
        "    def __init__(self, n_class):\n",
        "        super().__init__()\n",
        "        \n",
        "        self.base_model = torchvision.models.resnet18(True)\n",
        "        self.base_layers = list(self.base_model.children())\n",
        "        self.layer1 = nn.Sequential(\n",
        "            nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False),\n",
        "            self.base_layers[1],\n",
        "            self.base_layers[2])\n",
        "        self.layer2 = nn.Sequential(*self.base_layers[3:5])\n",
        "        self.layer3 = self.base_layers[5]\n",
        "        self.layer4 = self.base_layers[6]\n",
        "        self.layer5 = self.base_layers[7]\n",
        "        self.decode4 = Decoder(512, 256+256, 256)\n",
        "        self.decode3 = Decoder(256, 256+128, 256)\n",
        "        self.decode2 = Decoder(256, 128+64, 128)\n",
        "        self.decode1 = Decoder(128, 64+64, 64)\n",
        "        self.decode0 = nn.Sequential(\n",
        "            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),\n",
        "            nn.Conv2d(64, 32, kernel_size=3, padding=1, bias=False),\n",
        "            nn.Conv2d(32, 64, kernel_size=3, padding=1, bias=False)\n",
        "            )\n",
        "        self.conv_last = nn.Conv2d(64, n_class, 1)\n",
        "\n",
        "    def forward(self, input):\n",
        "        e1 = self.layer1(input) # 64,128,128\n",
        "        e2 = self.layer2(e1) # 64,64,64\n",
        "        e3 = self.layer3(e2) # 128,32,32\n",
        "        e4 = self.layer4(e3) # 256,16,16\n",
        "        f = self.layer5(e4) # 512,8,8\n",
        "        d4 = self.decode4(f, e4) # 256,16,16\n",
        "        d3 = self.decode3(d4, e3) # 256,32,32\n",
        "        d2 = self.decode2(d3, e2) # 128,64,64\n",
        "        d1 = self.decode1(d2, e1) # 64,128,128\n",
        "        d0 = self.decode0(d1) # 64,256,256\n",
        "        out = self.conv_last(d0) # 1,256,256\n",
        "        return out"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8YJji98EGsZU",
        "colab_type": "text"
      },
      "source": [
        "# Helper Functions"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1OVx5-R6hMKc",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train(train_loader, model):\n",
        "  running_loss = 0.0\n",
        "  data_size = len(train_data)\n",
        "\n",
        "  model.train()\n",
        "\n",
        "  for inputs, masks, labels in train_loader:\n",
        "    inputs, masks, labels = inputs.to(device), masks.to(device), labels.to(device)\n",
        "    optimizer.zero_grad()\n",
        "\n",
        "    with torch.set_grad_enabled(True):\n",
        "      logit = model(inputs)\n",
        "      loss = nn.BCEWithLogitsLoss()(logit.squeeze(1), masks.squeeze(1))\n",
        "      loss.backward()\n",
        "      optimizer.step()\n",
        "\n",
        "    running_loss += loss.item() * inputs.size(0)\n",
        "\n",
        "  epoch_loss = running_loss / data_size\n",
        "  return epoch_loss\n",
        "\n",
        "def test(test_loader, model):\n",
        "  running_loss = 0.0\n",
        "  data_size = len(test_loader)\n",
        "  predicts = []\n",
        "  truths = []\n",
        "\n",
        "  model.eval()\n",
        "\n",
        "  for inputs, masks in test_loader:\n",
        "    inputs, masks = inputs.to(device), masks.to(device)\n",
        "\n",
        "    with torch.set_grad_enabled(False):\n",
        "      outputs = model(inputs)\n",
        "      outputs = outputs[:, :, pad_left:pad_left + fine_size, pad_left:pad_left + fine_size].contiguous()\n",
        "      loss = nn.BCEWithLogitsLoss()(outputs.squeeze(1), masks.squeeze(1))\n",
        "\n",
        "    predicts.append(torch.sigmoid(outputs).detach().cpu().numpy()) \n",
        "    truths.append(masks.detach().cpu().numpy())\n",
        "    running_loss += loss.item() * inputs.size(0)\n",
        "\n",
        "  predicts = np.concatenate(predicts).squeeze()\n",
        "  truths = np.concatenate(truths).squeeze()\n",
        "  precision, _, _ = do_kaggle_metric(predicts, truths, 0.5)\n",
        "  precision = precision.mean()\n",
        "  epoch_loss = running_loss / data_size\n",
        "  return epoch_loss, precision\n",
        "\n",
        "def rle_encode(im):\n",
        "    '''\n",
        "    im: numpy array, 1 - mask, 0 - background\n",
        "    Returns run length as string formated\n",
        "    '''\n",
        "    pixels = im.flatten(order = 'F')\n",
        "    pixels = np.concatenate([[0], pixels, [0]])\n",
        "    runs = np.where(pixels[1:] != pixels[:-1])[0] + 1\n",
        "    runs[1::2] -= runs[::2]\n",
        "    return ' '.join(str(x) for x in runs)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KecZvuwl5Oey",
        "colab_type": "text"
      },
      "source": [
        "# Train"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GHbyq7iK3iRK",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "outputId": "b80547f2-d53f-43d9-950e-5b2c2da975d1"
      },
      "source": [
        "all_id = train_df['id'].values\n",
        "fold = []\n",
        "for i in range(5):\n",
        "  fold.append(train_df.loc[train_df['fold']==i, 'id'].values)\n",
        "\n",
        "salt = UNet(1)\n",
        "salt.to(device)"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "UNet(\n",
              "  (base_model): ResNet(\n",
              "    (conv1): Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
              "    (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "    (relu): ReLU(inplace=True)\n",
              "    (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
              "    (layer1): Sequential(\n",
              "      (0): BasicBlock(\n",
              "        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      )\n",
              "      (1): BasicBlock(\n",
              "        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      )\n",
              "    )\n",
              "    (layer2): Sequential(\n",
              "      (0): BasicBlock(\n",
              "        (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (downsample): Sequential(\n",
              "          (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
              "          (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        )\n",
              "      )\n",
              "      (1): BasicBlock(\n",
              "        (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "        (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      )\n",
              "    )\n",
              "    (layer3): Sequential(\n",
              "      (0): BasicBlock(\n",
              "        (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (downsample): Sequential(\n",
              "          (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
              "          (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        )\n",
              "      )\n",
              "      (1): BasicBlock(\n",
              "        (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "        (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      )\n",
              "    )\n",
              "    (layer4): Sequential(\n",
              "      (0): BasicBlock(\n",
              "        (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (downsample): Sequential(\n",
              "          (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
              "          (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        )\n",
              "      )\n",
              "      (1): BasicBlock(\n",
              "        (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "        (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      )\n",
              "    )\n",
              "    (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
              "    (fc): Linear(in_features=512, out_features=1000, bias=True)\n",
              "  )\n",
              "  (layer1): Sequential(\n",
              "    (0): Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)\n",
              "    (1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "    (2): ReLU(inplace=True)\n",
              "  )\n",
              "  (layer2): Sequential(\n",
              "    (0): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
              "    (1): Sequential(\n",
              "      (0): BasicBlock(\n",
              "        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      )\n",
              "      (1): BasicBlock(\n",
              "        (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "        (relu): ReLU(inplace=True)\n",
              "        (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "        (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      )\n",
              "    )\n",
              "  )\n",
              "  (layer3): Sequential(\n",
              "    (0): BasicBlock(\n",
              "      (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
              "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      (relu): ReLU(inplace=True)\n",
              "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      (downsample): Sequential(\n",
              "        (0): Conv2d(64, 128, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
              "        (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      )\n",
              "    )\n",
              "    (1): BasicBlock(\n",
              "      (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      (relu): ReLU(inplace=True)\n",
              "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "    )\n",
              "  )\n",
              "  (layer4): Sequential(\n",
              "    (0): BasicBlock(\n",
              "      (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
              "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      (relu): ReLU(inplace=True)\n",
              "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      (downsample): Sequential(\n",
              "        (0): Conv2d(128, 256, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
              "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      )\n",
              "    )\n",
              "    (1): BasicBlock(\n",
              "      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      (relu): ReLU(inplace=True)\n",
              "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "    )\n",
              "  )\n",
              "  (layer5): Sequential(\n",
              "    (0): BasicBlock(\n",
              "      (conv1): Conv2d(256, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
              "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      (relu): ReLU(inplace=True)\n",
              "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      (downsample): Sequential(\n",
              "        (0): Conv2d(256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False)\n",
              "        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      )\n",
              "    )\n",
              "    (1): BasicBlock(\n",
              "      (conv1): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "      (relu): ReLU(inplace=True)\n",
              "      (conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
              "    )\n",
              "  )\n",
              "  (decode4): Decoder(\n",
              "    (up): ConvTranspose2d(512, 256, kernel_size=(2, 2), stride=(2, 2))\n",
              "    (conv_relu): Sequential(\n",
              "      (0): Conv2d(512, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "      (1): ReLU(inplace=True)\n",
              "    )\n",
              "  )\n",
              "  (decode3): Decoder(\n",
              "    (up): ConvTranspose2d(256, 256, kernel_size=(2, 2), stride=(2, 2))\n",
              "    (conv_relu): Sequential(\n",
              "      (0): Conv2d(384, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "      (1): ReLU(inplace=True)\n",
              "    )\n",
              "  )\n",
              "  (decode2): Decoder(\n",
              "    (up): ConvTranspose2d(256, 128, kernel_size=(2, 2), stride=(2, 2))\n",
              "    (conv_relu): Sequential(\n",
              "      (0): Conv2d(192, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "      (1): ReLU(inplace=True)\n",
              "    )\n",
              "  )\n",
              "  (decode1): Decoder(\n",
              "    (up): ConvTranspose2d(128, 64, kernel_size=(2, 2), stride=(2, 2))\n",
              "    (conv_relu): Sequential(\n",
              "      (0): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
              "      (1): ReLU(inplace=True)\n",
              "    )\n",
              "  )\n",
              "  (decode0): Sequential(\n",
              "    (0): Upsample(scale_factor=2.0, mode=bilinear)\n",
              "    (1): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "    (2): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
              "  )\n",
              "  (conv_last): Conv2d(64, 1, kernel_size=(1, 1), stride=(1, 1))\n",
              ")"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "09E37d074u2K",
        "colab_type": "code",
        "outputId": "7927cb6e-7845-4f63-da0b-c07385fec9f0",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        }
      },
      "source": [
        "for idx in range(5):\n",
        "\n",
        "  if idx == 1:\n",
        "    break\n",
        "\n",
        "  # Setup optimizer\n",
        "  scheduler_step = epoch // snapshot\n",
        "  optimizer = torch.optim.SGD(salt.parameters(), lr=max_lr, momentum=momentum, weight_decay=weight_decay)\n",
        "  lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, scheduler_step, min_lr)\n",
        "\n",
        "  # Load data\n",
        "  train_id = np.setdiff1d(all_id, fold[idx])\n",
        "  val_id = fold[idx]\n",
        "\n",
        "  X_train, y_train = trainImageFetch(train_id)\n",
        "  X_val, y_val = trainImageFetch(val_id)\n",
        "\n",
        "  train_data = SaltDataset(X_train, 'train', y_train, pad_left=27, pad_right=27)\n",
        "  val_data = SaltDataset(X_val, 'val', y_val, pad_left=27, pad_right=27)\n",
        "\n",
        "  train_loader = DataLoader(train_data,\n",
        "                            shuffle=RandomSampler(train_data), \n",
        "                            batch_size=batch_size) \n",
        "\n",
        "  val_loader = DataLoader(val_data,\n",
        "                            shuffle=False, \n",
        "                            batch_size=batch_size) \n",
        "\n",
        "  num_snapshot = 0\n",
        "  best_acc = 0\n",
        "\n",
        "  for epoch_ in range(epoch):\n",
        "    train_loss = train(train_loader, salt)\n",
        "    val_loss, accuracy = test(val_loader, salt)\n",
        "    lr_scheduler.step()\n",
        "\n",
        "    if accuracy > best_acc:\n",
        "      best_acc = accuracy\n",
        "      best_param = salt.state_dict()\n",
        "\n",
        "    if (epoch_ + 1) % scheduler_step == 0:\n",
        "      torch.save(best_param, save_weight + weight_name + str(idx) + str(num_snapshot) + '.pth')\n",
        "      optimizer = torch.optim.SGD(salt.parameters(), lr=max_lr, momentum=momentum, weight_decay=weight_decay)\n",
        "      lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, scheduler_step, min_lr)\n",
        "      num_snapshot += 1\n",
        "      best_acc = 0\n",
        "\n",
        "    print('epoch: {} train_loss: {:.3f} val_loss: {:.3f} val_accuracy: {:.3f}'.format(epoch_ + 1, train_loss, val_loss, accuracy))"
      ],
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 3190/3190 [00:01<00:00, 2349.81it/s]\n",
            "100%|██████████| 810/810 [00:00<00:00, 2232.46it/s]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "epoch: 1 train_loss: 0.552 val_loss: 12.026 val_accuracy: 0.385\n",
            "epoch: 2 train_loss: 0.469 val_loss: 7.533 val_accuracy: 0.216\n",
            "epoch: 3 train_loss: 0.424 val_loss: 8.726 val_accuracy: 0.213\n",
            "epoch: 4 train_loss: 0.407 val_loss: 6.621 val_accuracy: 0.352\n",
            "epoch: 5 train_loss: 0.383 val_loss: 6.971 val_accuracy: 0.272\n",
            "epoch: 6 train_loss: 0.388 val_loss: 6.317 val_accuracy: 0.384\n",
            "epoch: 7 train_loss: 0.359 val_loss: 10.014 val_accuracy: 0.365\n",
            "epoch: 8 train_loss: 0.326 val_loss: 7.225 val_accuracy: 0.331\n",
            "epoch: 9 train_loss: 0.300 val_loss: 5.103 val_accuracy: 0.570\n",
            "epoch: 10 train_loss: 0.289 val_loss: 5.256 val_accuracy: 0.468\n",
            "epoch: 11 train_loss: 0.250 val_loss: 3.828 val_accuracy: 0.625\n",
            "epoch: 12 train_loss: 0.217 val_loss: 3.834 val_accuracy: 0.600\n",
            "epoch: 13 train_loss: 0.184 val_loss: 3.655 val_accuracy: 0.666\n",
            "epoch: 14 train_loss: 0.168 val_loss: 2.817 val_accuracy: 0.708\n",
            "epoch: 15 train_loss: 0.153 val_loss: 3.231 val_accuracy: 0.628\n",
            "epoch: 16 train_loss: 0.140 val_loss: 3.385 val_accuracy: 0.704\n",
            "epoch: 17 train_loss: 0.131 val_loss: 2.925 val_accuracy: 0.694\n",
            "epoch: 18 train_loss: 0.103 val_loss: 3.025 val_accuracy: 0.716\n",
            "epoch: 19 train_loss: 0.094 val_loss: 3.106 val_accuracy: 0.690\n",
            "epoch: 20 train_loss: 0.097 val_loss: 2.986 val_accuracy: 0.716\n",
            "epoch: 21 train_loss: 0.071 val_loss: 3.395 val_accuracy: 0.732\n",
            "epoch: 22 train_loss: 0.065 val_loss: 3.367 val_accuracy: 0.730\n",
            "epoch: 23 train_loss: 0.055 val_loss: 3.949 val_accuracy: 0.736\n",
            "epoch: 24 train_loss: 0.054 val_loss: 3.847 val_accuracy: 0.727\n",
            "epoch: 25 train_loss: 0.042 val_loss: 3.943 val_accuracy: 0.735\n",
            "epoch: 26 train_loss: 0.039 val_loss: 4.714 val_accuracy: 0.728\n",
            "epoch: 27 train_loss: 0.049 val_loss: 3.623 val_accuracy: 0.720\n",
            "epoch: 28 train_loss: 0.036 val_loss: 4.172 val_accuracy: 0.733\n",
            "epoch: 29 train_loss: 0.035 val_loss: 3.569 val_accuracy: 0.739\n",
            "epoch: 30 train_loss: 0.033 val_loss: 3.791 val_accuracy: 0.731\n",
            "epoch: 31 train_loss: 0.030 val_loss: 4.319 val_accuracy: 0.741\n",
            "epoch: 32 train_loss: 0.029 val_loss: 4.459 val_accuracy: 0.745\n",
            "epoch: 33 train_loss: 0.026 val_loss: 4.232 val_accuracy: 0.740\n",
            "epoch: 34 train_loss: 0.025 val_loss: 4.516 val_accuracy: 0.743\n",
            "epoch: 35 train_loss: 0.024 val_loss: 4.534 val_accuracy: 0.743\n",
            "epoch: 36 train_loss: 0.023 val_loss: 4.487 val_accuracy: 0.739\n",
            "epoch: 37 train_loss: 0.023 val_loss: 4.650 val_accuracy: 0.736\n",
            "epoch: 38 train_loss: 0.021 val_loss: 4.722 val_accuracy: 0.736\n",
            "epoch: 39 train_loss: 0.021 val_loss: 4.842 val_accuracy: 0.744\n",
            "epoch: 40 train_loss: 0.021 val_loss: 4.894 val_accuracy: 0.740\n",
            "epoch: 41 train_loss: 0.020 val_loss: 4.798 val_accuracy: 0.744\n",
            "epoch: 42 train_loss: 0.019 val_loss: 5.068 val_accuracy: 0.742\n",
            "epoch: 43 train_loss: 0.019 val_loss: 5.027 val_accuracy: 0.741\n",
            "epoch: 44 train_loss: 0.020 val_loss: 4.882 val_accuracy: 0.736\n",
            "epoch: 45 train_loss: 0.019 val_loss: 5.069 val_accuracy: 0.743\n",
            "epoch: 46 train_loss: 0.018 val_loss: 5.285 val_accuracy: 0.740\n",
            "epoch: 47 train_loss: 0.019 val_loss: 5.013 val_accuracy: 0.730\n",
            "epoch: 48 train_loss: 0.019 val_loss: 5.160 val_accuracy: 0.738\n",
            "epoch: 49 train_loss: 0.019 val_loss: 5.081 val_accuracy: 0.736\n",
            "epoch: 50 train_loss: 0.018 val_loss: 4.897 val_accuracy: 0.738\n",
            "epoch: 51 train_loss: 0.024 val_loss: 5.067 val_accuracy: 0.726\n",
            "epoch: 52 train_loss: 0.046 val_loss: 4.947 val_accuracy: 0.730\n",
            "epoch: 53 train_loss: 0.078 val_loss: 4.905 val_accuracy: 0.690\n",
            "epoch: 54 train_loss: 0.054 val_loss: 3.931 val_accuracy: 0.710\n",
            "epoch: 55 train_loss: 0.038 val_loss: 3.638 val_accuracy: 0.725\n",
            "epoch: 56 train_loss: 0.051 val_loss: 3.868 val_accuracy: 0.735\n",
            "epoch: 57 train_loss: 0.037 val_loss: 4.303 val_accuracy: 0.727\n",
            "epoch: 58 train_loss: 0.034 val_loss: 4.453 val_accuracy: 0.746\n",
            "epoch: 59 train_loss: 0.028 val_loss: 4.706 val_accuracy: 0.747\n",
            "epoch: 60 train_loss: 0.025 val_loss: 5.331 val_accuracy: 0.757\n",
            "epoch: 61 train_loss: 0.029 val_loss: 4.535 val_accuracy: 0.739\n",
            "epoch: 62 train_loss: 0.025 val_loss: 4.385 val_accuracy: 0.741\n",
            "epoch: 63 train_loss: 0.021 val_loss: 4.829 val_accuracy: 0.743\n",
            "epoch: 64 train_loss: 0.019 val_loss: 5.204 val_accuracy: 0.753\n",
            "epoch: 65 train_loss: 0.018 val_loss: 4.938 val_accuracy: 0.750\n",
            "epoch: 66 train_loss: 0.022 val_loss: 5.006 val_accuracy: 0.748\n",
            "epoch: 67 train_loss: 0.020 val_loss: 5.112 val_accuracy: 0.747\n",
            "epoch: 68 train_loss: 0.020 val_loss: 4.719 val_accuracy: 0.751\n",
            "epoch: 69 train_loss: 0.017 val_loss: 5.432 val_accuracy: 0.763\n",
            "epoch: 70 train_loss: 0.016 val_loss: 5.503 val_accuracy: 0.756\n",
            "epoch: 71 train_loss: 0.015 val_loss: 5.089 val_accuracy: 0.747\n",
            "epoch: 72 train_loss: 0.015 val_loss: 5.294 val_accuracy: 0.757\n",
            "epoch: 73 train_loss: 0.014 val_loss: 5.441 val_accuracy: 0.753\n",
            "epoch: 74 train_loss: 0.014 val_loss: 5.743 val_accuracy: 0.745\n",
            "epoch: 75 train_loss: 0.014 val_loss: 6.027 val_accuracy: 0.750\n",
            "epoch: 76 train_loss: 0.015 val_loss: 5.152 val_accuracy: 0.739\n",
            "epoch: 77 train_loss: 0.014 val_loss: 5.414 val_accuracy: 0.747\n",
            "epoch: 78 train_loss: 0.013 val_loss: 5.402 val_accuracy: 0.743\n",
            "epoch: 79 train_loss: 0.013 val_loss: 5.756 val_accuracy: 0.745\n",
            "epoch: 80 train_loss: 0.013 val_loss: 5.849 val_accuracy: 0.747\n",
            "epoch: 81 train_loss: 0.012 val_loss: 5.857 val_accuracy: 0.742\n",
            "epoch: 82 train_loss: 0.012 val_loss: 5.760 val_accuracy: 0.740\n",
            "epoch: 83 train_loss: 0.012 val_loss: 5.862 val_accuracy: 0.746\n",
            "epoch: 84 train_loss: 0.012 val_loss: 6.369 val_accuracy: 0.747\n",
            "epoch: 85 train_loss: 0.014 val_loss: 5.819 val_accuracy: 0.742\n",
            "epoch: 86 train_loss: 0.012 val_loss: 5.870 val_accuracy: 0.752\n",
            "epoch: 87 train_loss: 0.013 val_loss: 5.734 val_accuracy: 0.746\n",
            "epoch: 88 train_loss: 0.012 val_loss: 5.868 val_accuracy: 0.745\n",
            "epoch: 89 train_loss: 0.012 val_loss: 5.962 val_accuracy: 0.744\n",
            "epoch: 90 train_loss: 0.011 val_loss: 6.109 val_accuracy: 0.745\n",
            "epoch: 91 train_loss: 0.011 val_loss: 6.075 val_accuracy: 0.740\n",
            "epoch: 92 train_loss: 0.011 val_loss: 6.420 val_accuracy: 0.748\n",
            "epoch: 93 train_loss: 0.011 val_loss: 6.329 val_accuracy: 0.750\n",
            "epoch: 94 train_loss: 0.011 val_loss: 6.571 val_accuracy: 0.746\n",
            "epoch: 95 train_loss: 0.012 val_loss: 6.080 val_accuracy: 0.744\n",
            "epoch: 96 train_loss: 0.011 val_loss: 6.083 val_accuracy: 0.743\n",
            "epoch: 97 train_loss: 0.011 val_loss: 6.297 val_accuracy: 0.750\n",
            "epoch: 98 train_loss: 0.011 val_loss: 6.198 val_accuracy: 0.744\n",
            "epoch: 99 train_loss: 0.011 val_loss: 6.452 val_accuracy: 0.750\n",
            "epoch: 100 train_loss: 0.011 val_loss: 6.336 val_accuracy: 0.746\n",
            "epoch: 101 train_loss: 0.012 val_loss: 6.072 val_accuracy: 0.754\n",
            "epoch: 102 train_loss: 0.014 val_loss: 5.850 val_accuracy: 0.742\n",
            "epoch: 103 train_loss: 0.014 val_loss: 5.686 val_accuracy: 0.744\n",
            "epoch: 104 train_loss: 0.013 val_loss: 6.137 val_accuracy: 0.744\n",
            "epoch: 105 train_loss: 0.013 val_loss: 5.907 val_accuracy: 0.744\n",
            "epoch: 106 train_loss: 0.012 val_loss: 5.990 val_accuracy: 0.744\n",
            "epoch: 107 train_loss: 0.013 val_loss: 5.951 val_accuracy: 0.744\n",
            "epoch: 108 train_loss: 0.082 val_loss: 4.679 val_accuracy: 0.688\n",
            "epoch: 109 train_loss: 0.030 val_loss: 4.559 val_accuracy: 0.755\n",
            "epoch: 110 train_loss: 0.037 val_loss: 4.128 val_accuracy: 0.690\n",
            "epoch: 111 train_loss: 0.023 val_loss: 4.901 val_accuracy: 0.759\n",
            "epoch: 112 train_loss: 0.026 val_loss: 4.040 val_accuracy: 0.753\n",
            "epoch: 113 train_loss: 0.018 val_loss: 4.822 val_accuracy: 0.731\n",
            "epoch: 114 train_loss: 0.016 val_loss: 5.465 val_accuracy: 0.762\n",
            "epoch: 115 train_loss: 0.014 val_loss: 4.905 val_accuracy: 0.755\n",
            "epoch: 116 train_loss: 0.018 val_loss: 5.394 val_accuracy: 0.752\n",
            "epoch: 117 train_loss: 0.014 val_loss: 5.432 val_accuracy: 0.750\n",
            "epoch: 118 train_loss: 0.013 val_loss: 5.361 val_accuracy: 0.748\n",
            "epoch: 119 train_loss: 0.016 val_loss: 5.404 val_accuracy: 0.756\n",
            "epoch: 120 train_loss: 0.013 val_loss: 5.650 val_accuracy: 0.769\n",
            "epoch: 121 train_loss: 0.012 val_loss: 5.753 val_accuracy: 0.754\n",
            "epoch: 122 train_loss: 0.012 val_loss: 5.976 val_accuracy: 0.748\n",
            "epoch: 123 train_loss: 0.011 val_loss: 5.682 val_accuracy: 0.760\n",
            "epoch: 124 train_loss: 0.011 val_loss: 5.824 val_accuracy: 0.756\n",
            "epoch: 125 train_loss: 0.011 val_loss: 6.260 val_accuracy: 0.754\n",
            "epoch: 126 train_loss: 0.011 val_loss: 6.045 val_accuracy: 0.749\n",
            "epoch: 127 train_loss: 0.010 val_loss: 6.412 val_accuracy: 0.756\n",
            "epoch: 128 train_loss: 0.010 val_loss: 6.201 val_accuracy: 0.752\n",
            "epoch: 129 train_loss: 0.010 val_loss: 6.597 val_accuracy: 0.757\n",
            "epoch: 130 train_loss: 0.010 val_loss: 6.909 val_accuracy: 0.752\n",
            "epoch: 131 train_loss: 0.010 val_loss: 6.407 val_accuracy: 0.750\n",
            "epoch: 132 train_loss: 0.010 val_loss: 6.700 val_accuracy: 0.755\n",
            "epoch: 133 train_loss: 0.009 val_loss: 6.862 val_accuracy: 0.754\n",
            "epoch: 134 train_loss: 0.009 val_loss: 7.022 val_accuracy: 0.752\n",
            "epoch: 135 train_loss: 0.009 val_loss: 6.680 val_accuracy: 0.747\n",
            "epoch: 136 train_loss: 0.009 val_loss: 6.922 val_accuracy: 0.757\n",
            "epoch: 137 train_loss: 0.009 val_loss: 6.675 val_accuracy: 0.748\n",
            "epoch: 138 train_loss: 0.009 val_loss: 6.862 val_accuracy: 0.751\n",
            "epoch: 139 train_loss: 0.010 val_loss: 6.274 val_accuracy: 0.739\n",
            "epoch: 140 train_loss: 0.009 val_loss: 6.516 val_accuracy: 0.740\n",
            "epoch: 141 train_loss: 0.009 val_loss: 6.610 val_accuracy: 0.747\n",
            "epoch: 142 train_loss: 0.009 val_loss: 6.947 val_accuracy: 0.748\n",
            "epoch: 143 train_loss: 0.009 val_loss: 6.871 val_accuracy: 0.747\n",
            "epoch: 144 train_loss: 0.009 val_loss: 6.854 val_accuracy: 0.749\n",
            "epoch: 145 train_loss: 0.009 val_loss: 6.546 val_accuracy: 0.741\n",
            "epoch: 146 train_loss: 0.009 val_loss: 6.768 val_accuracy: 0.741\n",
            "epoch: 147 train_loss: 0.009 val_loss: 7.304 val_accuracy: 0.746\n",
            "epoch: 148 train_loss: 0.009 val_loss: 6.938 val_accuracy: 0.744\n",
            "epoch: 149 train_loss: 0.009 val_loss: 7.188 val_accuracy: 0.752\n",
            "epoch: 150 train_loss: 0.010 val_loss: 6.699 val_accuracy: 0.743\n",
            "epoch: 151 train_loss: 0.010 val_loss: 7.024 val_accuracy: 0.748\n",
            "epoch: 152 train_loss: 0.011 val_loss: 6.108 val_accuracy: 0.742\n",
            "epoch: 153 train_loss: 0.011 val_loss: 5.876 val_accuracy: 0.741\n",
            "epoch: 154 train_loss: 0.010 val_loss: 6.132 val_accuracy: 0.749\n",
            "epoch: 155 train_loss: 0.010 val_loss: 6.254 val_accuracy: 0.746\n",
            "epoch: 156 train_loss: 0.010 val_loss: 6.537 val_accuracy: 0.754\n",
            "epoch: 157 train_loss: 0.010 val_loss: 6.547 val_accuracy: 0.757\n",
            "epoch: 158 train_loss: 0.010 val_loss: 6.978 val_accuracy: 0.755\n",
            "epoch: 159 train_loss: 0.010 val_loss: 6.334 val_accuracy: 0.755\n",
            "epoch: 160 train_loss: 0.009 val_loss: 6.595 val_accuracy: 0.750\n",
            "epoch: 161 train_loss: 0.009 val_loss: 6.590 val_accuracy: 0.742\n",
            "epoch: 162 train_loss: 0.009 val_loss: 6.489 val_accuracy: 0.748\n",
            "epoch: 163 train_loss: 0.011 val_loss: 6.162 val_accuracy: 0.750\n",
            "epoch: 164 train_loss: 0.009 val_loss: 6.429 val_accuracy: 0.744\n",
            "epoch: 165 train_loss: 0.009 val_loss: 6.273 val_accuracy: 0.746\n",
            "epoch: 166 train_loss: 0.009 val_loss: 6.344 val_accuracy: 0.745\n",
            "epoch: 167 train_loss: 0.009 val_loss: 6.564 val_accuracy: 0.746\n",
            "epoch: 168 train_loss: 0.009 val_loss: 6.636 val_accuracy: 0.745\n",
            "epoch: 169 train_loss: 0.009 val_loss: 6.653 val_accuracy: 0.747\n",
            "epoch: 170 train_loss: 0.009 val_loss: 6.868 val_accuracy: 0.749\n",
            "epoch: 171 train_loss: 0.009 val_loss: 6.994 val_accuracy: 0.738\n",
            "epoch: 172 train_loss: 0.010 val_loss: 9.573 val_accuracy: 0.612\n",
            "epoch: 173 train_loss: 0.127 val_loss: 3.590 val_accuracy: 0.712\n",
            "epoch: 174 train_loss: 0.034 val_loss: 4.624 val_accuracy: 0.741\n",
            "epoch: 175 train_loss: 0.025 val_loss: 4.870 val_accuracy: 0.714\n",
            "epoch: 176 train_loss: 0.016 val_loss: 5.152 val_accuracy: 0.730\n",
            "epoch: 177 train_loss: 0.013 val_loss: 5.365 val_accuracy: 0.744\n",
            "epoch: 178 train_loss: 0.012 val_loss: 5.956 val_accuracy: 0.739\n",
            "epoch: 179 train_loss: 0.011 val_loss: 6.129 val_accuracy: 0.744\n",
            "epoch: 180 train_loss: 0.011 val_loss: 6.015 val_accuracy: 0.733\n",
            "epoch: 181 train_loss: 0.010 val_loss: 6.236 val_accuracy: 0.737\n",
            "epoch: 182 train_loss: 0.010 val_loss: 6.263 val_accuracy: 0.739\n",
            "epoch: 183 train_loss: 0.010 val_loss: 6.228 val_accuracy: 0.740\n",
            "epoch: 184 train_loss: 0.010 val_loss: 6.437 val_accuracy: 0.739\n",
            "epoch: 185 train_loss: 0.009 val_loss: 6.812 val_accuracy: 0.753\n",
            "epoch: 186 train_loss: 0.009 val_loss: 6.591 val_accuracy: 0.743\n",
            "epoch: 187 train_loss: 0.011 val_loss: 8.279 val_accuracy: 0.727\n",
            "epoch: 188 train_loss: 0.037 val_loss: 4.981 val_accuracy: 0.733\n",
            "epoch: 189 train_loss: 0.014 val_loss: 5.505 val_accuracy: 0.736\n",
            "epoch: 190 train_loss: 0.012 val_loss: 5.459 val_accuracy: 0.731\n",
            "epoch: 191 train_loss: 0.011 val_loss: 5.868 val_accuracy: 0.732\n",
            "epoch: 192 train_loss: 0.010 val_loss: 6.162 val_accuracy: 0.738\n",
            "epoch: 193 train_loss: 0.010 val_loss: 6.027 val_accuracy: 0.734\n",
            "epoch: 194 train_loss: 0.010 val_loss: 6.136 val_accuracy: 0.738\n",
            "epoch: 195 train_loss: 0.010 val_loss: 6.361 val_accuracy: 0.737\n",
            "epoch: 196 train_loss: 0.010 val_loss: 6.370 val_accuracy: 0.739\n",
            "epoch: 197 train_loss: 0.010 val_loss: 6.417 val_accuracy: 0.744\n",
            "epoch: 198 train_loss: 0.010 val_loss: 6.210 val_accuracy: 0.731\n",
            "epoch: 199 train_loss: 0.010 val_loss: 6.530 val_accuracy: 0.741\n",
            "epoch: 200 train_loss: 0.009 val_loss: 6.556 val_accuracy: 0.740\n",
            "epoch: 201 train_loss: 0.010 val_loss: 6.660 val_accuracy: 0.734\n",
            "epoch: 202 train_loss: 0.012 val_loss: 5.425 val_accuracy: 0.737\n",
            "epoch: 203 train_loss: 0.013 val_loss: 15.187 val_accuracy: 0.431\n",
            "epoch: 204 train_loss: 0.216 val_loss: 2.894 val_accuracy: 0.698\n",
            "epoch: 205 train_loss: 0.112 val_loss: 3.105 val_accuracy: 0.708\n",
            "epoch: 206 train_loss: 0.060 val_loss: 3.373 val_accuracy: 0.703\n",
            "epoch: 207 train_loss: 0.046 val_loss: 3.968 val_accuracy: 0.711\n",
            "epoch: 208 train_loss: 0.029 val_loss: 4.243 val_accuracy: 0.664\n",
            "epoch: 209 train_loss: 0.071 val_loss: 4.891 val_accuracy: 0.633\n",
            "epoch: 210 train_loss: 0.031 val_loss: 4.412 val_accuracy: 0.742\n",
            "epoch: 211 train_loss: 0.030 val_loss: 4.274 val_accuracy: 0.736\n",
            "epoch: 212 train_loss: 0.019 val_loss: 4.516 val_accuracy: 0.755\n",
            "epoch: 213 train_loss: 0.016 val_loss: 4.947 val_accuracy: 0.758\n",
            "epoch: 214 train_loss: 0.014 val_loss: 5.265 val_accuracy: 0.755\n",
            "epoch: 215 train_loss: 0.013 val_loss: 5.269 val_accuracy: 0.750\n",
            "epoch: 216 train_loss: 0.012 val_loss: 5.582 val_accuracy: 0.762\n",
            "epoch: 217 train_loss: 0.011 val_loss: 5.715 val_accuracy: 0.757\n",
            "epoch: 218 train_loss: 0.016 val_loss: 5.453 val_accuracy: 0.758\n",
            "epoch: 219 train_loss: 0.012 val_loss: 5.915 val_accuracy: 0.758\n",
            "epoch: 220 train_loss: 0.011 val_loss: 6.113 val_accuracy: 0.766\n",
            "epoch: 221 train_loss: 0.010 val_loss: 6.118 val_accuracy: 0.760\n",
            "epoch: 222 train_loss: 0.010 val_loss: 6.428 val_accuracy: 0.757\n",
            "epoch: 223 train_loss: 0.010 val_loss: 6.467 val_accuracy: 0.760\n",
            "epoch: 224 train_loss: 0.010 val_loss: 6.414 val_accuracy: 0.756\n",
            "epoch: 225 train_loss: 0.010 val_loss: 6.608 val_accuracy: 0.760\n",
            "epoch: 226 train_loss: 0.010 val_loss: 6.285 val_accuracy: 0.753\n",
            "epoch: 227 train_loss: 0.009 val_loss: 6.214 val_accuracy: 0.748\n",
            "epoch: 228 train_loss: 0.009 val_loss: 6.597 val_accuracy: 0.756\n",
            "epoch: 229 train_loss: 0.009 val_loss: 6.819 val_accuracy: 0.752\n",
            "epoch: 230 train_loss: 0.009 val_loss: 6.855 val_accuracy: 0.755\n",
            "epoch: 231 train_loss: 0.009 val_loss: 6.728 val_accuracy: 0.758\n",
            "epoch: 232 train_loss: 0.009 val_loss: 6.979 val_accuracy: 0.754\n",
            "epoch: 233 train_loss: 0.009 val_loss: 6.588 val_accuracy: 0.752\n",
            "epoch: 234 train_loss: 0.008 val_loss: 6.845 val_accuracy: 0.756\n",
            "epoch: 235 train_loss: 0.008 val_loss: 6.930 val_accuracy: 0.750\n",
            "epoch: 236 train_loss: 0.008 val_loss: 7.197 val_accuracy: 0.756\n",
            "epoch: 237 train_loss: 0.008 val_loss: 7.059 val_accuracy: 0.755\n",
            "epoch: 238 train_loss: 0.008 val_loss: 7.108 val_accuracy: 0.753\n",
            "epoch: 239 train_loss: 0.008 val_loss: 7.029 val_accuracy: 0.751\n",
            "epoch: 240 train_loss: 0.008 val_loss: 6.982 val_accuracy: 0.755\n",
            "epoch: 241 train_loss: 0.008 val_loss: 6.931 val_accuracy: 0.752\n",
            "epoch: 242 train_loss: 0.008 val_loss: 6.974 val_accuracy: 0.755\n",
            "epoch: 243 train_loss: 0.008 val_loss: 6.849 val_accuracy: 0.753\n",
            "epoch: 244 train_loss: 0.008 val_loss: 7.219 val_accuracy: 0.756\n",
            "epoch: 245 train_loss: 0.008 val_loss: 7.065 val_accuracy: 0.755\n",
            "epoch: 246 train_loss: 0.008 val_loss: 7.124 val_accuracy: 0.755\n",
            "epoch: 247 train_loss: 0.008 val_loss: 7.198 val_accuracy: 0.753\n",
            "epoch: 248 train_loss: 0.008 val_loss: 7.272 val_accuracy: 0.755\n",
            "epoch: 249 train_loss: 0.008 val_loss: 7.184 val_accuracy: 0.755\n",
            "epoch: 250 train_loss: 0.008 val_loss: 7.501 val_accuracy: 0.753\n",
            "epoch: 251 train_loss: 0.008 val_loss: 7.428 val_accuracy: 0.755\n",
            "epoch: 252 train_loss: 0.009 val_loss: 6.426 val_accuracy: 0.744\n",
            "epoch: 253 train_loss: 0.018 val_loss: 5.804 val_accuracy: 0.745\n",
            "epoch: 254 train_loss: 0.011 val_loss: 5.467 val_accuracy: 0.743\n",
            "epoch: 255 train_loss: 0.010 val_loss: 6.244 val_accuracy: 0.751\n",
            "epoch: 256 train_loss: 0.009 val_loss: 6.185 val_accuracy: 0.742\n",
            "epoch: 257 train_loss: 0.009 val_loss: 6.383 val_accuracy: 0.750\n",
            "epoch: 258 train_loss: 0.009 val_loss: 6.549 val_accuracy: 0.746\n",
            "epoch: 259 train_loss: 0.009 val_loss: 6.318 val_accuracy: 0.747\n",
            "epoch: 260 train_loss: 0.008 val_loss: 6.613 val_accuracy: 0.753\n",
            "epoch: 261 train_loss: 0.008 val_loss: 6.576 val_accuracy: 0.749\n",
            "epoch: 262 train_loss: 0.008 val_loss: 6.859 val_accuracy: 0.748\n",
            "epoch: 263 train_loss: 0.009 val_loss: 6.903 val_accuracy: 0.659\n",
            "epoch: 264 train_loss: 0.086 val_loss: 4.316 val_accuracy: 0.728\n",
            "epoch: 265 train_loss: 0.040 val_loss: 7.545 val_accuracy: 0.706\n",
            "epoch: 266 train_loss: 0.022 val_loss: 4.808 val_accuracy: 0.746\n",
            "epoch: 267 train_loss: 0.013 val_loss: 5.315 val_accuracy: 0.751\n",
            "epoch: 268 train_loss: 0.011 val_loss: 5.948 val_accuracy: 0.758\n",
            "epoch: 269 train_loss: 0.010 val_loss: 5.989 val_accuracy: 0.754\n",
            "epoch: 270 train_loss: 0.009 val_loss: 6.307 val_accuracy: 0.753\n",
            "epoch: 271 train_loss: 0.009 val_loss: 6.521 val_accuracy: 0.758\n",
            "epoch: 272 train_loss: 0.009 val_loss: 6.621 val_accuracy: 0.758\n",
            "epoch: 273 train_loss: 0.009 val_loss: 6.594 val_accuracy: 0.757\n",
            "epoch: 274 train_loss: 0.008 val_loss: 6.539 val_accuracy: 0.754\n",
            "epoch: 275 train_loss: 0.008 val_loss: 6.757 val_accuracy: 0.760\n",
            "epoch: 276 train_loss: 0.008 val_loss: 6.690 val_accuracy: 0.748\n",
            "epoch: 277 train_loss: 0.008 val_loss: 6.682 val_accuracy: 0.750\n",
            "epoch: 278 train_loss: 0.008 val_loss: 6.924 val_accuracy: 0.755\n",
            "epoch: 279 train_loss: 0.008 val_loss: 7.235 val_accuracy: 0.755\n",
            "epoch: 280 train_loss: 0.008 val_loss: 7.101 val_accuracy: 0.754\n",
            "epoch: 281 train_loss: 0.014 val_loss: 6.218 val_accuracy: 0.759\n",
            "epoch: 282 train_loss: 0.009 val_loss: 6.925 val_accuracy: 0.759\n",
            "epoch: 283 train_loss: 0.008 val_loss: 6.738 val_accuracy: 0.757\n",
            "epoch: 284 train_loss: 0.008 val_loss: 6.920 val_accuracy: 0.761\n",
            "epoch: 285 train_loss: 0.008 val_loss: 7.308 val_accuracy: 0.763\n",
            "epoch: 286 train_loss: 0.008 val_loss: 7.000 val_accuracy: 0.761\n",
            "epoch: 287 train_loss: 0.008 val_loss: 7.028 val_accuracy: 0.759\n",
            "epoch: 288 train_loss: 0.008 val_loss: 7.351 val_accuracy: 0.761\n",
            "epoch: 289 train_loss: 0.008 val_loss: 6.848 val_accuracy: 0.759\n",
            "epoch: 290 train_loss: 0.008 val_loss: 7.020 val_accuracy: 0.764\n",
            "epoch: 291 train_loss: 0.007 val_loss: 7.200 val_accuracy: 0.763\n",
            "epoch: 292 train_loss: 0.007 val_loss: 7.082 val_accuracy: 0.760\n",
            "epoch: 293 train_loss: 0.007 val_loss: 7.354 val_accuracy: 0.761\n",
            "epoch: 294 train_loss: 0.007 val_loss: 7.303 val_accuracy: 0.760\n",
            "epoch: 295 train_loss: 0.007 val_loss: 7.010 val_accuracy: 0.755\n",
            "epoch: 296 train_loss: 0.007 val_loss: 7.099 val_accuracy: 0.760\n",
            "epoch: 297 train_loss: 0.007 val_loss: 7.051 val_accuracy: 0.751\n",
            "epoch: 298 train_loss: 0.007 val_loss: 7.169 val_accuracy: 0.761\n",
            "epoch: 299 train_loss: 0.007 val_loss: 7.535 val_accuracy: 0.762\n",
            "epoch: 300 train_loss: 0.007 val_loss: 7.101 val_accuracy: 0.754\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NgQ8uhBTr6Ts",
        "colab_type": "text"
      },
      "source": [
        "# Test"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bFhhsn5pxbkY",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 202
        },
        "outputId": "0e475b8b-93a2-4faa-be73-1322863612b4"
      },
      "source": [
        "test_id = [x[:-4] for x in os.listdir(test_image_dir) if x[-4:] == '.png']\n",
        "image_test = testImageFetch(test_id)\n",
        "overall_pred_101 = np.zeros((len(test_id), 101, 101), dtype=np.float32)\n",
        "\n",
        "for step in range(1, 6):\n",
        "\n",
        "  print('Predicting Snapshot', step)\n",
        "  pred_null = []\n",
        "\n",
        "  # Load weight\n",
        "  param = torch.load(save_weight + weight_name + '0' + str(step) + '.pth')\n",
        "  salt.load_state_dict(param)\n",
        "\n",
        "  # Dataloader\n",
        "  test_data = SaltDataset(image_test, mode='test', fine_size=fine_size, pad_left=pad_left, pad_right=pad_right)\n",
        "  test_loader = DataLoader(test_data,\n",
        "                            shuffle=False,\n",
        "                            batch_size=batch_size)\n",
        "  \n",
        "  # Prediction\n",
        "  salt.eval()\n",
        "  for images in tqdm(test_loader, total=len(test_loader)):\n",
        "    images = images.to(device)\n",
        "    with torch.set_grad_enabled(False):\n",
        "      pred = salt(images)\n",
        "      pred = torch.sigmoid(pred).squeeze(1).cpu().numpy()\n",
        "      pred = pred[:, pad_left:pad_left + fine_size, pad_left:pad_left + fine_size]\n",
        "      pred_null.append(pred)\n",
        "  \n",
        "  idx = 0\n",
        "  for i in range(len(pred_null)):\n",
        "    for j in range(batch_size):\n",
        "      overall_pred_101[idx] += cv2.resize(pred_null[i][j], dsize=(101, 101))\n",
        "      idx += 1"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 18000/18000 [00:13<00:00, 1320.24it/s]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Predicting Snapshot 1\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 1000/1000 [01:34<00:00, 10.38it/s]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Predicting Snapshot 2\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 1000/1000 [01:35<00:00, 10.47it/s]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Predicting Snapshot 3\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 1000/1000 [01:35<00:00, 10.46it/s]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Predicting Snapshot 4\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 1000/1000 [01:35<00:00, 10.42it/s]\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "Predicting Snapshot 5\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "100%|██████████| 1000/1000 [01:35<00:00, 10.55it/s]\n"
          ],
          "name": "stderr"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wkOT1GWjmdOC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "submission = pd.DataFrame({'id':test_id, 'rle_mask':list(overall_pred_101)})\n",
        "submission['rle_mask'] = submission['rle_mask'].map(lambda x: rle_encode(x>5*0.5))\n",
        "submission.set_index('id', inplace=True)\n",
        "\n",
        "sample_submission = pd.read_csv('tgs-salt-identification-challenge/sample_submission.csv')\n",
        "sample_submission.set_index('id', inplace=True)\n",
        "submission = submission.reindex(sample_submission.index)\n",
        "submission.reset_index(inplace=True)\n",
        "submission.to_csv('submission.csv', index=False)"
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}